from fastapi import APIRouter, UploadFile
from fastapi.websockets import WebSocket
from fastapi.responses import JSONResponse, Response
import aiofiles
from .upload_record import *
from .chunk import *
from .models import FileInfo

router = APIRouter(tags=['File Upload'])


async def receive_chunk(websocket: WebSocket) -> bool:
    try:
        buffer = await websocket.receive_bytes()
        chunk = BufferChunk(buffer)
        record = UploadRecord.load(chunk.fid)
        res = ChunkRecord(fid=chunk.fid, offset=chunk.offset)
        if not record:
            res.state = ChunkState.Rejected
            res.message = 'Chunk not allowed, post file information first!'
            print('Chunk Rejectd:',chunk.fid)
            await websocket.send_json(res.json())
            return False
        if chunk.ended:
            print('End chunk:', res.json())
            record.complete()
            res.state = ChunkState.Ended
            res.message = record.file_name
            await websocket.send_json(res.json())
            return False

        await record.save_chunk(chunk)
        res.state = ChunkState.Accepted
        print('Chunk Saved:', res.json())
        await websocket.send_json(res.json())
        return True
    except Exception as e:
        res.state = ChunkState.Failed
        res.message = str(e)
        print('Save Chunk Failed:', res.json())
        await websocket.send_json(res.json())
        return False


@router.websocket("/ws/upload")
async def websocket(websocket: WebSocket):
    await websocket.accept()
    while True:
        go = await receive_chunk(websocket)
        if not go:
            break
    await websocket.close()


@router.websocket("/ws/check")
async def check(websocket: WebSocket):
    await websocket.accept()
    while True:
        try:
            res = await websocket.receive_text()
            print(res)
        except:
            break
    await websocket.close()


@router.post('/files/upload', response_class=JSONResponse)
async def request_upload(file_info: FileInfo):
    print('File Info:',file_info)
    file_name = file_info.file_name.strip()
    record = UploadRecord.make_record(file_name,file_info.file_size)
    if record.is_complete and file_info.reupload:
        record = UploadRecord.make_record(file_name,file_info.file_size,True)

    return record.json()


@router.post('/files/upload/small', response_class=JSONResponse, description="upload samll files")
async def upload_small(file: UploadFile):
    try:
        file_name = file.filename.strip()
        file_path = upload_dir/file_name
        async with aiofiles.open(file_path, 'wb') as f:
            data = await file.read()
            await f.write(data)
            print('Saved small file:', file_name)
        res = {
            "file_path": str(file_path)
        }
        return JSONResponse(res)
    except Exception as e:
        res = {
            "file_path": str(file_path),
            "error": str(e)
        }
        return JSONResponse(res, 411)
